#include "Scanner.h"
#include "Utils.h"
#include "FileDetails.h"
#include "StlFile.h"
#include "ReportJob.h"
#include "GroupFinishedJob.h"
#include "ScanGroup.h"
#include "StreamScanGroup.h"
#include "StreamGroupFinishedJob.h"
#include "ClamAVScanJob.h"
#include <Windows.h>
#include <exception>

namespace GPUCC
{
	Scanner::Scanner(const Config& Conf, bool CpuOnly) : m_ScannedFilesCounter(0), m_CpuOnly(CpuOnly), m_Config(Conf), m_Engine(*this), m_GroupsLimits(Conf.GetMaxNumberOfGroups() * Conf.GetMaxSizeOfGroup()), m_StreamGroupsLimits(Conf.GetMaxNumberOfStreamGroups() * Conf.GetMaxSizeOfStreamGroup())
	{
	}

	Scanner::~Scanner()
	{
		this->m_Groups.clear();
		this->m_StreamGroups.clear();
	}

	bool Scanner::Init()
	{
		bool Success = false;

		Success = this->m_ClamAV.Init(this->m_Config.GetSignaturesFolder());
		
		if (Success)
		{
			if (!this->m_CpuOnly)
			{
				Success = this->InitGPU();
			}
		}

		return (Success);
	}

	void Scanner::Cleanup()
	{
		// todo: implement!!!!
	}

	std::vector<InfectedDetails> Scanner::ScanDrives(bool ScanRemovableDrives)
	{
		char DriveLettersBuffer[26 * 4 + 1];
		ZeroMemory(DriveLettersBuffer, sizeof(DriveLettersBuffer));

		if (this->StartThreads())
		{
			if (GetLogicalDriveStrings(sizeof(DriveLettersBuffer) - 1, DriveLettersBuffer))
			{
				char* CurrentDrive = DriveLettersBuffer;
			
				do 
				{
					UINT CurrentDriveType = GetDriveType(CurrentDrive);

					if ((CurrentDriveType == DRIVE_FIXED) ||
						((CurrentDriveType == DRIVE_REMOVABLE) && (ScanRemovableDrives)))
					{
						std::string DriveName(CurrentDrive);

						DriveName.reserve(MAX_PATH);

						DriveName.erase(DriveName.length() - 1);

						this->ScanSingleFolder(DriveName);
					}
				} while ((CurrentDrive += 4) && (*CurrentDrive));
			}

			this->CompleteScan();
		}

		return (this->m_Results);
	}

	std::vector<InfectedDetails> Scanner::ScanFolder(std::string& Path, bool Recursive)
	{
		if (this->StartThreads())
		{
			Path.reserve(MAX_PATH);

			this->ScanSingleFolder(Path, Recursive);

			this->CompleteScan();
		}

		return (this->m_Results);
	}

	unsigned int Scanner::GetScannedFilesCount()
	{
		return (this->m_ScannedFilesCounter);
	}

	bool Scanner::InitGPU()
	{
		bool Success = false;

		std::vector<std::string> SignatureFiles = Utils::GetAllFilesNames(this->m_Config.GetSignaturesFolder(), "*.hdb");
		std::vector<std::string> UnwantedProgramSignatureFiles = Utils::GetAllFilesNames(this->m_Config.GetSignaturesFolder(), "*.hdu");

		SignatureFiles.insert(SignatureFiles.end(), UnwantedProgramSignatureFiles.begin(), UnwantedProgramSignatureFiles.end());

		for (std::vector<std::string>::const_iterator CurrentFileName = SignatureFiles.begin(); CurrentFileName != SignatureFiles.end(); ++CurrentFileName)
		{
			std::string CurrentFilePath(this->m_Config.GetSignaturesFolder());

			CurrentFilePath.append("\\");
			CurrentFilePath.append(*CurrentFileName);
			CurrentFilePath = Utils::GetFullPath(CurrentFilePath);

			std::vector<MD5Signature> Signatures;
				
			try
			{
				Signatures = MD5Signature::ParseSignaturesFile(CurrentFilePath);
			}
			catch (std::exception& ex)
			{
				// todo: log error???
			}

			this->m_Signatures.insert(this->m_Signatures.end(), Signatures.begin(), Signatures.end());
		}
			
		Success = this->m_Signatures.size() != 0;

		if (Success)
		{
			Success = this->InitOpenCL();

			if (Success)
			{
				this->m_Groups.reserve(this->m_Config.GetMaxNumberOfGroups());

				for (unsigned int GroupCount = 0; GroupCount < this->m_Config.GetMaxNumberOfGroups(); ++GroupCount)
				{
					this->m_Groups.push_back(std::make_shared<ScanGroup>(this->m_Config.GetMaxSizeOfGroup(),
																		this->m_Config.GetMaxFileSizeWithoutStreaming(),
																		(KernelHandler&)this->m_Engine,
																		&this->m_Program,
																		&this->m_Context,
																		&this->m_CommandQueue[GroupCount]));
				}

				this->m_StreamGroups.reserve(this->m_Config.GetMaxNumberOfStreamGroups());

				for (unsigned int GroupCount = 0; GroupCount < this->m_Config.GetMaxNumberOfStreamGroups(); ++GroupCount)
				{
					 this->m_StreamGroups.push_back(std::make_shared<StreamScanGroup>(this->m_Config.GetMaxSizeOfStreamGroup(),
																					this->m_Config.GetGPUMaxFileSize(),
																					(KernelHandler&)this->m_Engine,
																					&this->m_Program,
																					&this->m_Context,
																					&this->m_StreamCommandQueue[GroupCount]));
				}
			}
		}

		return (Success);
	}

	bool Scanner::InitOpenCL()
	{
		bool Success = false;

		std::vector<cl::Platform> Platforms;

		if (cl::Platform::get(&Platforms) == CL_SUCCESS)
		{
			std::string Value;
			std::vector<cl::Device> GPUDevices;

			for (std::vector<cl::Platform>::iterator CurrentPlatform = Platforms.begin(); CurrentPlatform != Platforms.end(); ++CurrentPlatform)
			{
				CurrentPlatform->getInfo(CL_PLATFORM_NAME, &Value);

				if (this->m_Config.GetPlatformName().compare(Value) == 0)
				{
					this->m_Platform = *CurrentPlatform;

					if (CurrentPlatform->getDevices(CL_DEVICE_TYPE_GPU, &GPUDevices) == CL_SUCCESS)
					{
						for (std::vector<cl::Device>::iterator CurrentDevice = GPUDevices.begin(); CurrentDevice != GPUDevices.end(); ++CurrentDevice)
						{
							CurrentDevice->getInfo(CL_DEVICE_NAME, &Value);

							if (this->m_Config.GetGPUName().compare(Value) == 0)
							{
								this->m_GPU = *CurrentDevice;

								Success = true;
							}
						}
					}
				}
			}
		}

		if (Success)
		{
			cl_int Result;
			cl_context_properties Properties[3];

			Success = false;
			
			Properties[0] = CL_CONTEXT_PLATFORM;
			Properties[1] = (cl_context_properties)(this->m_Platform)();
			Properties[2] = 0;

			this->m_Context = cl::Context(this->m_GPU, Properties, NULL, NULL, &Result);

			if (Result == CL_SUCCESS)
			{
				VECTOR_CLASS<cl::Device> dev;
				VECTOR_CLASS<cl_context_properties> props;
				this->m_Context.getInfo(CL_CONTEXT_DEVICES, &dev);
				this->m_Context.getInfo(CL_CONTEXT_PROPERTIES, &props);

				std::vector<std::string> strings;
				cl::Program::Sources sources;

				if (this->LoadSources(strings, sources))
				{
					this->m_Program = cl::Program(this->m_Context, sources, &Result);

					if (Result == CL_SUCCESS)
					{
						Result = this->m_Program.build(NULL, NULL, NULL);
						
						if (Result == CL_SUCCESS)
						{
							for (int i = 0; i < this->m_Config.GetMaxNumberOfGroups(); ++i)
							{
								// todo: decide how much queue are required! signle glocal queue\queue per group\queue per number of groups etc..
								this->m_CommandQueue.push_back(cl::CommandQueue(this->m_Context, this->m_GPU, 0, &Result));
							}

							if (Result == CL_SUCCESS)
							{
								for (int i = 0; i < this->m_Config.GetMaxNumberOfStreamGroups(); ++i)
								{
									// todo: decide how much queue are required! signle glocal queue\queue per group\queue per number of groups etc..
									this->m_StreamCommandQueue.push_back(cl::CommandQueue(this->m_Context, this->m_GPU, 0, &Result));
								}

								if (Result == CL_SUCCESS)
								{
									Success = true;
								}
							}
						}
					}
				}
			}
		}

		return (Success);
	}

	bool Scanner::LoadSources(std::vector<std::string>& strings, cl::Program::Sources& sources)
	{
		bool Result = true;
		std::vector<StlFile> Files;

		Files.push_back(Utils::GetFullPath("MD5Hash.cl"));

		for (std::vector<StlFile>::iterator CurrentFile = Files.begin(); ((CurrentFile != Files.end()) && (Result)); ++CurrentFile)
		{
			if (!CurrentFile->Open())
			{
				Result = false;
			}
			else
			{
				strings.push_back(std::string(std::istreambuf_iterator<char>(CurrentFile->GetStream()), (std::istreambuf_iterator<char>())));
				CurrentFile->Close();

				sources.push_back(std::make_pair(strings[strings.size() - 1].c_str(), strings[strings.size() - 1].length() + 1));
			}
		}

		return (Result);
	}

	bool Scanner::StartThreads()
	{
		bool Success = false;

		this->m_Results.clear();
		this->m_Results.shrink_to_fit();

		if (this->m_Engine.Start())
		{
			unsigned int NumberOfThreads = 1;

			if (this->m_CpuOnly)
			{
				SYSTEM_INFO sysinfo;

				GetSystemInfo(&sysinfo);

				NumberOfThreads = sysinfo.dwNumberOfProcessors;
			}

			Success = this->m_ClamAVThreads.Start(NumberOfThreads);

			if (!Success)
			{
				this->m_Engine.Stop();
			}
		}

		return (Success);
	}
	
	void Scanner::CompleteScan()
	{
		this->m_ClamAVThreads.StopAndWait();

		this->m_GroupsCS.Enter();

		for (unsigned int CurrentGroup = 0; CurrentGroup < this->m_Groups.size(); ++CurrentGroup)
		{
			if ((!this->m_Groups[CurrentGroup]->IsFull()) && (this->m_Groups[CurrentGroup]->GetItemsCount() > 0))
			{
				this->m_Groups[CurrentGroup]->Scan();
			}
		}

		this->m_GroupsCS.Leave();

		this->m_StreamGroupsCS.Enter();

		for (unsigned int CurrentGroup = 0; CurrentGroup < this->m_StreamGroups.size(); ++CurrentGroup)
		{
			if ((!this->m_StreamGroups[CurrentGroup]->IsFull()) && (this->m_StreamGroups[CurrentGroup]->GetItemsCount() > 0))
			{
				this->m_StreamGroups[CurrentGroup]->Scan();
			}
		}

		this->m_StreamGroupsCS.Leave();

		for (unsigned int CurrentGroup = 0; CurrentGroup < this->m_GroupsLimits.GetCount(); ++CurrentGroup)
		{
			this->m_GroupsLimits.Get();
		}

		this->m_GroupsLimits.Release(this->m_GroupsLimits.GetCount());

		for (unsigned int CurrentGroup = 0; CurrentGroup < this->m_StreamGroupsLimits.GetCount(); ++CurrentGroup)
		{
			this->m_StreamGroupsLimits.Get();
		}

		this->m_StreamGroupsLimits.Release(this->m_StreamGroupsLimits.GetCount());

		this->m_Engine.Stop();
	}

	void Scanner::ScanSingleFolder(std::string& Path, bool Recursive)
	{
		WIN32_FIND_DATA ffd;
		HANDLE hFind = INVALID_HANDLE_VALUE;

		Path.append("\\*");

		hFind = FindFirstFile(Path.c_str(), &ffd);

		Path.erase(Path.length() - 1);

		if (hFind != INVALID_HANDLE_VALUE)
		{
			std::string CurrentItem;

			do
			{
				CurrentItem.assign(ffd.cFileName);

				Path.append(CurrentItem);

				if (ffd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
				{
					if ((Recursive) && (CurrentItem.compare(".") != 0) && (CurrentItem.compare("..") != 0))
					{
						this->ScanSingleFolder(Path, Recursive);
					}
				}
				else
				{
					this->ScanFile(Path, ffd.nFileSizeLow, ffd.nFileSizeHigh);
				}

				Path.erase(Path.length() - CurrentItem.length());
			} while (FindNextFile(hFind, &ffd) != 0);

			FindClose(hFind);
		}

		Path.erase(Path.length() - 1);
	}

	void Scanner::ScanFile(const std::string& Path, const DWORD SizeLow, const DWORD SizeHigh)
	{
		LARGE_INTEGER FileSize;
		FileSize.LowPart = SizeLow;
		FileSize.HighPart = SizeHigh;

		this->ScanFile(Path, FileSize);
	}

	std::vector<unsigned char> ConvertBytesToHex(const std::string& Value);

	void Scanner::ScanFile(const std::string& Path, const LARGE_INTEGER& Size)
	{
		unsigned int FileSize = FileDetails::MAX_FILE_SIZE;
		bool StreamFile = false;

		if (Size.HighPart == 0)
		{
			FileSize = Size.LowPart;
		}

		if (FileSize == 0)
		{
			return;
		}

		if ((this->m_Config.IsMaxFileSizeLimit()) && (FileSize > this->m_Config.GetMaxFileSize()))
		{
			return;
		}

		++this->m_ScannedFilesCounter;

		if ((this->m_CpuOnly) || (FileSize > this->m_Config.GetGPUMaxFileSize()))
		{
			//this->m_ClamAVThreads.AddJobExecution(std::make_shared<ClamAVScanJob>(Path, this->m_ClamAV, this->m_Engine));
		}
		else if (FileSize > this->m_Config.GetMaxFileSizeWithoutStreaming())
		{
			ScanFileStreamed(Path, FileSize);
		}
		else
		{
			ScanEntireFile(Path, FileSize);
		}
	}

	void Scanner::ScanEntireFile(const std::string& Path, unsigned int FileSize)
	{
		std::shared_ptr<BaseScanGroup> group;
		
		this->m_GroupsLimits.Get();

		this->m_GroupsCS.Enter();

		for (unsigned int CurrentGroup = 0; CurrentGroup < this->m_Groups.size(); ++CurrentGroup)
		{
			if (!this->m_Groups[CurrentGroup]->IsFull())
			{
				group = this->m_Groups[CurrentGroup];

				break;
			}
		}

		this->m_GroupsCS.Leave();

		if (!group->AddItem(Path, FileSize))
		{
			this->m_GroupsLimits.Release();
		}
		else
		{
			if (group->IsFull())
			{
				group->Scan();
			}
		}
	}

	void Scanner::ScanFileStreamed(const std::string& Path, unsigned int FileSize)
	{
		std::shared_ptr<BaseScanGroup> group;
		
		this->m_StreamGroupsLimits.Get();

		this->m_StreamGroupsCS.Enter();

		for (unsigned int CurrentGroup = 0; CurrentGroup < this->m_StreamGroups.size(); ++CurrentGroup)
		{
			if (!this->m_StreamGroups[CurrentGroup]->IsFull())
			{
				group = this->m_StreamGroups[CurrentGroup];

				break;
			}
		}

		this->m_StreamGroupsCS.Leave();

		if (!group->AddItem(Path, FileSize))
		{
			this->m_StreamGroupsLimits.Release();
		}
		else
		{
			if (group->IsFull())
			{
				group->Scan();
			}
		}
	}

	Scanner::Engine::Engine(Scanner& Instance) : m_Scanner(Instance)
	{
	}

	Scanner::Engine::~Engine()
	{
	}

	bool Scanner::Engine::Start()
	{
		return (this->m_ThreadPool.Start(1));
	}

	void Scanner::Engine::Stop()
	{
		this->m_ThreadPool.StopAndWait();
	}

	void Scanner::Engine::OnInfected(const std::string& Path, const std::string& SignatureName)
	{
		this->m_ThreadPool.AddJobExecution(std::make_shared<ReportJob>(this->m_Scanner, Path, SignatureName));
	}

	void Scanner::Engine::OnKernelOutputRead(BaseScanGroup* Instance)
	{
		this->m_ThreadPool.AddJobExecution(std::make_shared<BaseScanGroup::CompareJob>(Instance, this->m_Scanner.m_Signatures, *this));
	}

	void Scanner::Engine::OnGroupFinished(BaseScanGroup* Instance)
	{
		this->m_ThreadPool.AddJobExecution(std::make_shared<GroupFinishedJob>(this->m_Scanner, Instance));
	}

	void Scanner::Engine::OnStreamGroupFinished(BaseScanGroup* Instance)
	{
		this->m_ThreadPool.AddJobExecution(std::make_shared<StreamGroupFinishedJob>(this->m_Scanner, Instance));
	}
}